'''
Author: Ryan
Email: 409266255@qq.com
Date: 2025-09-23 22:11:45

LastEditors: Ryan
LastEditTime: 2025-09-28 12:13:20
Description: efinance数据获取

PS：港美股接口
    - get_quote_history：获取股票的 K 线数据
    - get_realtime_quotes：获取单个或者多个市场行情的最新状况
    - get_latest_quote：获取沪深市场、港股、美股多只股票的实时涨幅情况（建议使用get_realtime_quotes）
    - API文档：https://efinance.readthedocs.io/en/latest/api.html

FilePath: /qopilot/qopilot/data/datasource/qop_efinance.py
Copyright (c) 2025 by Ryan, All Rights Reserved. 
'''

import efinance as ef
import pandas as pd
import logging
import os
from enum import Enum

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class MarketType(Enum):
    """市场类型"""
    US = '美股'                     # 美股市场行情
    US_CN = '中概股'                # 中国概念股市场行情
    HK = '港股'                     # 港股市场行情
    CN_SH = '沪深A股'               # 沪深A股市场行情
    CN_SH_A = '沪A'                # 沪市A股市场行情
    CN_SZ_A = '深A'                # 深市A股市场行情
    CN_BJ_A = '北A'                # 北证A股市场行情
    CN_CYB = '创业板'               # 创业板市场行情
    CN_KGP = '科创板'               # 科创板市场行情
    CN_HGK = '沪股通'               # 沪股通市场行情
    CN_SGK = '深股通'               # 深股通市场行情
    CN_XG = '新股'                  # 沪深新股市场行情
    CN_KZZ = '可转债'               # 沪深可转债市场行情
    CN_HYBK = '行业板块'            # 行业板块市场行情
    CN_GNBK = '概念板块'            # 概念板块市场行情
    CN_HS_INDEX = '沪深系列指数'    # 沪深系列指数市场行情
    CN_SH_INDEX = '上证系列指数'    # 上证系列指数市场行情
    CN_SZ_INDEX = '深证系列指数'    # 深证系列指数市场行情
    CN_FUTURES = '期货'             # 期货市场行情
    CN_ETF = 'ETF'                  # ETF 基金市场行情   
    CN_LOF = 'LOF'                  # LOF 基金市场行情


# 历史数据列配置
history_column_config = {
    'rename': {
        '股票代码': 'code', 
        '股票名称': 'name',
        '日期': 'date',
        '开盘': 'open',
        '收盘': 'close',
        '最高': 'high',
        '最低': 'low',
        '成交量': 'volume',
        '成交额': 'amount',
        '振幅': 'amplitude',
        '涨跌幅': 'pct_change',
        '涨跌额': 'change',
        '换手率': 'turnover_rate'
    },
    'keep': ['code', 'date', 'open', 'close', 'high', 'low', 'volume', 'amount', 'pct_change'],
    'drop': []
}

# 实时数据列配置
realtime_column_config = {
    'rename': {
        '股票代码': 'code',
        '股票名称': 'name',
        '今开': 'open',
        '最新价': 'close',
        '最高': 'high',
        '最低': 'low',
        '涨跌幅': 'pct_change',
        '涨跌额': 'change',
        '换手率': 'turnover_rate',
        '量比': 'turnover_rate',
        '动态市盈率': 'pe_ratio',
        '成交量': 'volume',
        '成交额': 'amount',
        '昨日收盘': 'pre_close',
        '总市值': 'total_market_value',
        '流通市值': 'traded_market_value',
        '行情ID': 'quote_id',
        '市场类型': 'market_type',
        '更新时间': 'update_time',
        '最新交易日': 'latest_trade_date'
    },
    'keep': ['code', 'date', 'open', 'close'],
    'drop': []
}


def standardize_column_names(df, column_config=None):
    """
    更灵活的列重命名和过滤函数

    Args:
        df: 输入的DataFrame
        column_config: 列配置字典，格式为：
            {
                'rename': {'原列名': '新列名'},  # 重命名映射
                'keep': ['列名1', '列名2'],     # 要保留的列（重命名后）
                'drop': ['列名1', '列名2']      # 要删除的列（重命名后）
            }

    Returns:
        DataFrame: 处理后的DataFrame
    """
    if column_config is None:
        return df

    df_result = df.copy()

    # 重命名列
    if 'rename' in column_config and column_config['rename']:
        df_result = df_result.rename(columns=column_config['rename'])

    # 删除列
    if 'drop' in column_config and column_config['drop']:
        df_result = df_result.drop(columns=column_config['drop'], errors='ignore')

    # 保留指定列
    if 'keep' in column_config and column_config['keep']:
        available_columns = [col for col in column_config['keep'] if col in df_result.columns]
        df_result = df_result[available_columns]

    return df_result


def test():
    """测试列重命名和过滤功能"""

    # stock_code = 'AAPL'
    # df = ef.stock.get_quote_history(stock_code, beg='20250101', end='20250924', fqt=2)
    # df_standardized = standardize_column_names(df, history_column_config)
    # print(df_standardized.tail(5))

    df = ef.stock.get_realtime_quotes(['港股'])
    # df = ef.stock.get_realtime_quotes(['创业板', '港股'])
    # ef.stock.get_realtime_quotes(['美股'])
    print(df)


class EFinanceData:
    def __init__(self):
        pass

    def get_basic(self, stock_market: list = [MarketType.HK]) -> pd.DataFrame:
        return ef.stock.get_realtime_quotes(stock_market)

    def get_quote_history(self, stock_code: str, beg: str, end: str, fqt: int) -> pd.DataFrame:
        pass

    def get_realtime_quotes(self, markets: list) -> pd.DataFrame:
        pass


class HKEFinanceData(EFinanceData):
    """港股EFinance数据获取类"""

    def __init__(self):
        logger.info("港股数据获取器初始化成功")

    def get_hk_basic(self) -> pd.DataFrame:
        """
        获取港股列表

        Returns:
            DataFrame: 港股列表数据（已过滤掉ts_code以'8'开头的数据）

        Api Ref:
            https://tushare.pro/document/2?doc_id=191
        """
        try:
            logger.info("开始获取港股列表")

            # 调用Tushare接口
            df = self.get_basic([MarketType.HK.value])

            # 过滤掉ts_code以'8'开头的数据
            # if not df.empty and 'ts_code' in df.columns:
            #     original_count = len(df)
            #     df = df[~df['ts_code'].str.startswith('8')]
            #     filtered_count = len(df)
            #     logger.info(f"过滤前共{original_count}条记录，过滤后剩余{filtered_count}条记录")
            # else:
            #     logger.info("数据为空或缺少ts_code字段，无需过滤")

            # logger.info(f"成功获取港股列表，共{len(df)}条记录")

            # 将df写入Excel文件
            # 确保目标目录存在
            # 使用相对路径，从当前文件位置（qopilot/data/datasource/）到项目根目录的database文件夹
            current_dir = os.path.dirname(os.path.abspath(__file__))
            project_root = os.path.join(current_dir, '..', '..', '..')
            database_dir = os.path.join(project_root, 'database')
            os.makedirs(database_dir, exist_ok=True)

            excel_filename = f"hk_stocks_{pd.Timestamp.now().strftime('%Y%m%d_%H%M%S')}.xlsx"
            excel_path = os.path.join(database_dir, excel_filename)
            df.to_excel(excel_path, index=False, engine='openpyxl')
            logger.info(f"港股数据已导出到Excel文件: {excel_path}")

            return df

        except Exception as e:
            logger.error(f"获取港股列表失败: {str(e)}")
            raise


class USEEFinanceData(EFinanceData):
    """美股EFinance数据获取类"""

    def __init__(self):
        self.ef = ef
        self.get_realtime_quotes = self.ef.stock.get_realtime_quotes()
        self.get_quote_history = self.ef.stock.get_quote_history()


if __name__ == "__main__":
    # test()  # 完整演示
    hk = HKEFinanceData()
    print(hk.get_hk_basic())
